import logging
import os.path as osp
import time
import math
from functools import partial

import numpy as np
import torch
import torch_geometric.transforms as T
from numpy.random import default_rng
from ogb.nodeproppred import PygNodePropPredDataset
from ogb.graphproppred import PygGraphPropPredDataset
from torch_geometric.datasets import (Amazon, Coauthor, GNNBenchmarkDataset, TUDataset,
                                      WikipediaNetwork, ZINC)
from torch_geometric.graphgym.config import cfg
from torch_geometric.graphgym.loader import load_pyg, load_ogb, set_dataset_attr
from torch_geometric.graphgym.register import register_loader

from graphgps.loader.planetoid import Planetoid
from graphgps.loader.dataset.aqsol_molecules import AQSOL
from graphgps.loader.dataset.coco_superpixels import COCOSuperpixels
from graphgps.loader.dataset.malnet_tiny import MalNetTiny
from graphgps.loader.dataset.voc_superpixels import VOCSuperpixels
from graphgps.loader.split_generator import (prepare_splits,
                                             set_dataset_splits)
from graphgps.transform.posenc_stats import compute_posenc_stats
from graphgps.transform.transforms import (pre_transform_in_memory,
                                           generate_splits, # not the same as split_generator above.
                                           typecast_x, concat_x_and_pos,
                                           clip_graphs_to_size, move_node_feat_to_x)
from graphgps.transform.expander_edges import generate_random_expander
from graphgps.transform.dist_transforms import (add_dist_features, add_reverse_edges,
                                                 add_self_loops, effective_resistances, 
                                                 effective_resistance_embedding,
                                                 effective_resistances_from_embedding)


def log_loaded_dataset(dataset, format, name):
    logging.info(f"[*] Loaded dataset '{name}' from '{format}':")
    logging.info(f"  {dataset.data}")
    # logging.info(f"  undirected: {dataset[0].is_undirected()}")
    # len(dataset): 打印数据集中图的总数量（多个图数据的集合）。
    logging.info(f"  num graphs: {len(dataset)}")

    # 计算数据集中总的节点数量。如果 dataset.data 包含 num_nodes 属性，就直接使用它。否则，通过节点特征矩阵 x 的第一维度（即节点数）来确定节点数量。
    total_num_nodes = 0
    if hasattr(dataset.data, 'num_nodes'):
        total_num_nodes = dataset.data.num_nodes
    elif hasattr(dataset.data, 'x'):
        total_num_nodes = dataset.data.x.size(0)
    # 计算每个图的平均节点数（总节点数除以图的数量）。
    logging.info(f"  avg num_nodes/graph: "
                 f"{total_num_nodes // len(dataset)}")
    # 分别打印节点特征数和边特征数。
    logging.info(f"  num node features: {dataset.num_node_features}")
    logging.info(f"  num edge features: {dataset.num_edge_features}")
    # 检查数据集是否有 num_tasks 属性，并记录任务数量（通常与多任务学习相关）。
    if hasattr(dataset, 'num_tasks'):
        logging.info(f"  num tasks: {dataset.num_tasks}")

    # 检查数据集中是否有标签 y
    if hasattr(dataset.data, 'y') and dataset.data.y is not None:
        # 如果标签为列表（如 ogbg-code2 数据集），则类别数不可用 (n/a)。
        if isinstance(dataset.data.y, list):
            # A special case for ogbg-code2 dataset.
            logging.info(f"  num classes: n/a")
        # 如果 y 是浮点数并且 y.numel() == y.size(0)，则判断这是回归任务（每个样本一个连续值）。
        elif dataset.data.y.numel() == dataset.data.y.size(0) and \
                torch.is_floating_point(dataset.data.y):
            logging.info(f"  num classes: (appears to be a regression task)")
        # 分类任务：否则，输出类别数 dataset.num_classes
        else:
            logging.info(f"  num classes: {dataset.num_classes}")
    # 针对边预测任务，检查 train_edge_label 或 edge_label 是否存在。这种任务可能是预测边的存在与否（分类）或边的权重（回归）。
    elif hasattr(dataset.data, 'train_edge_label') or hasattr(dataset.data, 'edge_label'):
        # Edge/link prediction task.
        if hasattr(dataset.data, 'train_edge_label'):
            labels = dataset.data.train_edge_label  # Transductive link task
        else:
            labels = dataset.data.edge_label  # Inductive link task
        # 如果标签是浮点数并且每条边有一个数值，则推断这是一个回归任务。
        if labels.numel() == labels.size(0) and \
                torch.is_floating_point(labels):
            logging.info(f"  num edge classes: (probably a regression task)")
        # 否则，通过 torch.unique 计算不同的边类别数。
        else:
            logging.info(f"  num edge classes: {len(torch.unique(labels))}")

    # 这部分代码被注释掉了，但它计算和打印了数据集中每个图的节点数分布（图的规模分布）。通过 np.histogram 生成分布直方图并打印每个区间的图数。
    ## Show distribution of graph sizes.
    # graph_sizes = [d.num_nodes if hasattr(d, 'num_nodes') else d.x.shape[0]
    #                for d in dataset]
    # hist, bin_edges = np.histogram(np.array(graph_sizes), bins=10)
    # logging.info(f'   Graph size distribution:')
    # logging.info(f'     mean: {np.mean(graph_sizes)}')
    # for i, (start, end) in enumerate(zip(bin_edges[:-1], bin_edges[1:])):
    #     logging.info(
    #         f'     bin {i}: [{start:.2f}, {end:.2f}]: '
    #         f'{hist[i]} ({hist[i] / hist.sum() * 100:.2f}%)'
    #     )


@register_loader('custom_master_loader')
def load_dataset_master(format, name, dataset_dir):
    """
    Master loader that controls loading of all datasets, overshadowing execution
    of any default GraphGym dataset loader. Default GraphGym dataset loader are
    instead called from this function, the format keywords `PyG` and `OGB` are
    reserved for these default GraphGym loaders.

    Custom transforms and dataset splitting is applied to each loaded dataset.

    Args:
        format: dataset format name that identifies Dataset class
        name: dataset name to select from the class identified by `format`
        dataset_dir: path where to store the processed dataset

    Returns:
        PyG dataset object with applied perturbation transforms and data splits
    """
    # 这段代码首先检查format是否以PyG-开头，如果是，则提取format中标识数据集的部分（pyg_dataset_id），然后将数据集的存储路径拼接为dataset_dir。
    if format.startswith('PyG-'):
        pyg_dataset_id = format.split('-', 1)[1]
        dataset_dir = osp.join(dataset_dir, pyg_dataset_id)

        #todo 需要修改：加载指定数据集
        if pyg_dataset_id == 'GNNBenchmarkDataset':
            #dataset_dir = '/datasets/GNN'
            # 调用自定义的preformat_GNNBenchmarkDataset函数来加载并预处理该数据集。
            dataset = preformat_GNNBenchmarkDataset(dataset_dir, name)

        elif pyg_dataset_id == 'MalNetTiny':
            #dataset_dir = '/datasets/GNN'
            # 调用preformat_MalNetTiny加载特定的MalNetTiny数据集，并通过feature_set设置特征。
            dataset = preformat_MalNetTiny(dataset_dir, feature_set=name)

        elif pyg_dataset_id == 'Amazon':
            # 加载Amazon数据集，特定数据集（photo或computers）会被进一步处理，添加数据集分割和反向边，并根据配置是否添加自环。
            dataset = Amazon(dataset_dir, name)
            if name == "photo" or name == "computers":
                pre_transform_in_memory(dataset, partial(generate_splits, g_split = cfg.dataset.split[0]))
                pre_transform_in_memory(dataset, partial(add_reverse_edges))
                if cfg.prep.add_self_loops:
                  pre_transform_in_memory(dataset, partial(add_self_loops))

        elif pyg_dataset_id == 'Coauthor':
            dataset = Coauthor(dataset_dir, name)
            if name == "physics" or name == "cs":
                pre_transform_in_memory(dataset, partial(generate_splits, g_split = cfg.dataset.split[0]))
                pre_transform_in_memory(dataset, partial(add_reverse_edges))
                if cfg.prep.add_self_loops:
                  pre_transform_in_memory(dataset, partial(add_self_loops))

        elif pyg_dataset_id == 'Planetoid':
            # dataset = Planetoid(dataset_dir, name)
            dataset = Planetoid(dataset_dir, name, split='random', train_percent= cfg.prep.train_percent)
            # dataset = Planetoid(dataset_dir, name, split='random', num_train_per_class = 4725, num_val = 1575, num_test = 1575)
            # Citeseer
            # dataset = Planetoid(dataset_dir, name, split='random', num_train_per_class = 1996, num_val = 665, num_test = 666)
            if name == "PubMed":
                pre_transform_in_memory(dataset, partial(typecast_x, type_str='float'))
            if cfg.prep.add_reverse_edges == True:
              pre_transform_in_memory(dataset, partial(add_reverse_edges))
            if cfg.prep.add_self_loops == True:
              pre_transform_in_memory(dataset, partial(add_self_loops))

        elif pyg_dataset_id == 'TUDataset':
            dataset = preformat_TUDataset(dataset_dir, name)

        elif pyg_dataset_id == 'VOCSuperpixels':
            #dataset_dir = '/datasets/LRGB'
            dataset = preformat_VOCSuperpixels(dataset_dir, name,
                                               cfg.dataset.slic_compactness)

        elif pyg_dataset_id == 'COCOSuperpixels':
            #dataset_dir = '/datasets/LRGB'
            dataset = preformat_COCOSuperpixels(dataset_dir, name,
                                                cfg.dataset.slic_compactness)


        elif pyg_dataset_id == 'WikipediaNetwork':
            if name == 'crocodile':
                raise NotImplementedError(f"crocodile not implemented yet")
            dataset = WikipediaNetwork(dataset_dir, name)

        elif pyg_dataset_id == 'ZINC':
            dataset = preformat_ZINC(dataset_dir, name)
            
        elif pyg_dataset_id == 'AQSOL':
            dataset = preformat_AQSOL(dataset_dir, name)

        else:
            raise ValueError(f"Unexpected PyG Dataset identifier: {format}")

    # GraphGym default loader for Pytorch Geometric datasets
    elif format == 'PyG':
        # 如果format为PyG，则使用GraphGym的默认加载器load_pyg来加载数据集。
        dataset = load_pyg(name, dataset_dir)

    elif format == 'OGB':
        if name.startswith('ogbg'):
            dataset = preformat_OGB_Graph(dataset_dir, name.replace('_', '-'))

        elif name.startswith('PCQM4Mv2-'):
            subset = name.split('-', 1)[1]
            dataset = preformat_OGB_PCQM4Mv2(dataset_dir, subset)
        elif name.startswith('ogbn'):
            dataset = preformat_ogbn(dataset_dir, name)

        elif name.startswith('peptides-'):
            #dataset_dir = '/datasets/LRGB'
            dataset = preformat_Peptides(dataset_dir, name)

        ### Link prediction datasets.
        elif name.startswith('ogbl-'):
            # GraphGym default loader.
            dataset = load_ogb(name, dataset_dir)
            # OGB link prediction datasets are binary classification tasks,
            # however the default loader creates float labels => convert to int.
            def convert_to_int(ds, prop):
                tmp = getattr(ds.data, prop).int()
                set_dataset_attr(ds, prop, tmp, len(tmp))
            convert_to_int(dataset, 'train_edge_label')
            convert_to_int(dataset, 'val_edge_label')
            convert_to_int(dataset, 'test_edge_label')

        elif name.startswith('PCQM4Mv2Contact-'):
            #dataset_dir = '/datasets/LRGB'
            dataset = preformat_PCQM4Mv2Contact(dataset_dir, name)

        else:
            raise ValueError(f"Unsupported OGB(-derived) dataset: {name}")
    else:
        raise ValueError(f"Unknown data format: {format}")

    # 使用自定义的log_loaded_dataset函数记录加载数据集的基本信息，包括图的数量、节点的平均数量、节点/边特征的数量等。
    log_loaded_dataset(dataset, format, name)

    # Precompute necessary statistics for positional encodings.
    # 检查配置文件中启用了哪些位置编码，并对启用的编码进行预计算（例如有效电阻的计算等）。
    # 初始化一个空列表 pe_enabled_list，用于存储启用了的 PE 类型名称（如 Laplacian PE, RW PE 等）。
    pe_enabled_list = []
    # 遍历配置项 cfg 中的每一个键值对
    for key, pecfg in cfg.items():
        # 如果键名以 posenc_ 开头且位置编码已启用 (pecfg.enable == True)，并且键名不是 posenc_ER（即排除 ER 位置编码），则：
        # 从键名中提取出 PE 类型名称（例如 posenc_RW 中的 RW）并添加到 pe_enabled_list 中。
        if key.startswith('posenc_') and pecfg.enable and (not key.startswith('posenc_ER')):
            pe_name = key.split('_', 1)[1]
            pe_enabled_list.append(pe_name)
            # 如果 pecfg 配置项中有 kernel 属性，并且 kernel 中定义了 times_func，则通过 eval() 函数执行这个函数并将其返回结果存储在 kernel.times 中。
            # 例如，某些 PE 可能基于特定的时间步长或频率参数来计算。
            # eval() 是一个 Python 内置函数，用于执行存储在字符串中的表达式。
            # 使用 logging.info 记录日志，打印 PE 名称和解析出的 kernel.times。
            if hasattr(pecfg, 'kernel'):
                # Generate kernel times if functional snippet is set.
                if pecfg.kernel.times_func:
                    pecfg.kernel.times = list(eval(pecfg.kernel.times_func))
                logging.info(f"Parsed {pe_name} PE kernel times / steps: "
                             f"{pecfg.kernel.times}")

    # 如果 pe_enabled_list 不为空（即存在启用的 PE），记录当前时间 start，并输出日志信息，告知正在计算 PE 的统计信息。
    # time.perf_counter() 是一种高精度的计时器，记录当前的时间戳，用于后续计算执行耗时。
    if pe_enabled_list:
        start = time.perf_counter()
        logging.info(f"Precomputing Positional Encoding statistics: "
                     f"{pe_enabled_list} for all graphs...")
        # Estimate directedness based on 10 graphs to save time.
        # 检查数据集中前10个图是否为无向图。
        # 调用每个图数据的 is_undirected() 方法，判断是否为无向图。
        # 使用 all() 函数，如果前10个图都是无向的，则 is_undirected 为 True，否则为 False。
        # 输出日志信息，记录估计的图的有向性。
        is_undirected = all(d.is_undirected() for d in dataset[:10])
        logging.info(f"  ...estimated to be undirected: {is_undirected}")
        #  使用 pre_transform_in_memory() 函数，对数据集中的每个图应用位置编码的统计信息计算。
        # 具体的计算通过 compute_posenc_stats 函数进行，其中：
        # pe_types=pe_enabled_list: 传入启用的 PE 类型列表。
        # is_undirected: 根据前面的估计结果，传入图的有向性。
        # cfg=cfg: 传入配置项 cfg 。
        # partial() 是 Python 中的 functools 模块，用于创建部分函数应用，将参数提前传递给函数。
        # show_progress=True: 显示进度条，提示处理过程的进展。
        pre_transform_in_memory(dataset,
                                partial(compute_posenc_stats,
                                        pe_types=pe_enabled_list,
                                        is_undirected=is_undirected,
                                        cfg=cfg),
                                show_progress=True
                                )
        # 计算整个 PE 统计计算过程的耗时：
        # 使用 time.perf_counter() 获取当前时间，减去之前记录的 start，得到耗时 elapsed。
        # 使用 time.strftime() 将耗时格式化为小时、分钟和秒，精确到小数点后两位。
        # 记录日志，输出 "Done! Took {timestr}"，表明 PE 计算完成及其耗时。
        elapsed = time.perf_counter() - start
        timestr = time.strftime('%H:%M:%S', time.gmtime(elapsed)) \
                  + f'{elapsed:.2f}'[-3:]
        logging.info(f"Done! Took {timestr}")

    # Other preprocessings:
    # adding expander edges:
    # 如果配置文件中要求添加扩展边，会调用generate_random_expander生成随机扩展边。
    if cfg.prep.exp:
        for j in range(cfg.prep.exp_count):
            start = time.perf_counter()
            logging.info(f"Adding expander edges (round {j}) ...")
            pre_transform_in_memory(dataset,
                                    partial(generate_random_expander,
                                            degree = cfg.prep.exp_deg,
                                            algorithm = cfg.prep.exp_algorithm,
                                            rng = None,
                                            max_num_iters = cfg.prep.exp_max_num_iters,
                                            exp_index = j),
                                    show_progress=True
                                    )
            elapsed = time.perf_counter() - start
            timestr = time.strftime('%H:%M:%S', time.gmtime(elapsed)) \
                      + f'{elapsed:.2f}'[-3:]
            logging.info(f"Done! Took {timestr}")


    # adding shortest path features
    # 根据配置文件，预先计算每个图的最短路径距离特征并将其添加到数据集中。
    if cfg.prep.dist_enable:
        start = time.perf_counter()
        logging.info(f"Precalculating node distances and shortest paths ...")
        is_undirected = dataset[0].is_undirected()
        Max_N = max([data.num_nodes for data in dataset])
        pre_transform_in_memory(dataset,
                                partial(add_dist_features,
                                        max_n = Max_N,
                                        is_undirected = is_undirected,
                                        cutoff = cfg.prep.dist_cutoff),
                                show_progress=True
                                )
        elapsed = time.perf_counter() - start
        timestr = time.strftime('%H:%M:%S', time.gmtime(elapsed)) \
                  + f'{elapsed:.2f}'[-3:]
        logging.info(f"Done! Took {timestr}")


    # adding effective resistance features
    # 根据配置文件，预计算图的有效电阻特征，并将其加入数据集中。MaxK是位置编码的维度，基于每个图的节点数量计算。
    if cfg.posenc_ERN.enable or cfg.posenc_ERE.enable:
        start = time.perf_counter()
        logging.info(f"Precalculating effective resistance for graphs ...")
        
        MaxK = max(
            [
                min(
                math.ceil(data.num_nodes//2), 
                math.ceil(8 * math.log(data.num_edges) / (cfg.posenc_ERN.accuracy**2))
                ) 
                for data in dataset
            ]
            )

        cfg.posenc_ERN.er_dim = MaxK
        logging.info(f"Choosing ER pos enc dim = {MaxK}")

        pre_transform_in_memory(dataset,
                                partial(effective_resistance_embedding,
                                        MaxK = MaxK,
                                        accuracy = cfg.posenc_ERN.accuracy,
                                        which_method = 0),
                                show_progress=True
                                )

        pre_transform_in_memory(dataset,
                        partial(effective_resistances_from_embedding,
                        normalize_per_node = False),
                        show_progress=True
                        )

        elapsed = time.perf_counter() - start
        timestr = time.strftime('%H:%M:%S', time.gmtime(elapsed)) \
                  + f'{elapsed:.2f}'[-3:]
        logging.info(f"Done! Took {timestr}")


    # This could not be done earlier because the training wants 'train_mask' etc.
    # Now after using gnn.head: inductive_node this is ok.
    if name == 'ogbn-arxiv' or name == 'ogbn-proteins':
      return dataset
    # Set standard dataset train/val/test splits
    if hasattr(dataset, 'split_idxs'):
        set_dataset_splits(dataset, dataset.split_idxs)
        delattr(dataset, 'split_idxs')

    # Verify or generate dataset train/val/test splits
    # 最后，生成或验证数据集的训练、验证和测试集分割。
    prepare_splits(dataset)

    # Precompute in-degree histogram if needed for PNAConv.
    # 检查配置文件中的 GNN 图卷积层类型 cfg.gt.layer_type 是否是 PNAConv，即通过 startswith('PNAConv') 来判断。如果是 PNAConv 类型的卷积层，同时 cfg.gt.pna_degrees 列表为空（len(cfg.gt.pna_degrees) == 0），就会执行接下来的操作。
    # PNAConv: 是一种用于图神经网络的卷积层，能够根据邻居节点的度数分布来增强聚合操作。为使其工作，通常需要输入节点的度数信息。
    # cfg.gt.pna_degrees: 用于存储图中每个节点的度数（或入度）的直方图。如果为空，则需要根据训练图计算这些度数信息。
    if cfg.gt.layer_type.startswith('PNAConv') and len(cfg.gt.pna_degrees) == 0:
        cfg.gt.pna_degrees = compute_indegree_histogram(
            dataset[dataset.data['train_graph_index']])

    return dataset


def compute_indegree_histogram(dataset):
    """Compute histogram of in-degree of nodes needed for PNAConv.

    Args:
        dataset: PyG Dataset object

    Returns:
        List where i-th value is the number of nodes with in-degree equal to `i`
    """
    from torch_geometric.utils import degree

    deg = torch.zeros(1000, dtype=torch.long)
    max_degree = 0
    for data in dataset:
        d = degree(data.edge_index[1],
                   num_nodes=data.num_nodes, dtype=torch.long)
        max_degree = max(max_degree, d.max().item())
        deg += torch.bincount(d, minlength=deg.numel())
    return deg.numpy().tolist()[:max_degree + 1]


def preformat_GNNBenchmarkDataset(dataset_dir, name):
    """Load and preformat datasets from PyG's GNNBenchmarkDataset.

    Args:
        dataset_dir: path where to store the cached dataset
        name: name of the specific dataset in the TUDataset class

    Returns:
        PyG dataset object
    """
    # 初始化一个空的列表 tf_list，用于存储对数据集应用的变换函数列表。
    tf_list = []
    # 检查 name 是否是 'MNIST' 或 'CIFAR10'，这两个数据集是支持的类型。如果是，则对数据集应用特定的变换。
    if name in ['MNIST', 'CIFAR10']:
        # 将 concat_x_and_pos 函数添加到 tf_list，它的作用是将特征值和坐标位置进行拼接。这个操作通常在图像数据中使用，以包含空间信息。
        tf_list = [concat_x_and_pos]  # concat pixel value and pos. coordinate
        # 将 typecast_x 函数的部分应用（partial function）添加到 tf_list 中，用于将数据类型强制转换为浮点数。
        # partial(typecast_x, type_str='float')：使用 functools.partial 创建一个部分函数，将 typecast_x 函数的 type_str 参数设定为 'float'。
        tf_list.append(partial(typecast_x, type_str='float'))
    else:
        ValueError(f"Loading dataset '{name}' from "
                   f"GNNBenchmarkDataset is not supported.")
    # 这段代码通过遍历训练集（train）、验证集（val）和测试集（test）来加载数据，并将它们合并为一个完整的数据集。
    # GNNBenchmarkDataset(root=dataset_dir, name=name, split=split)：从指定的目录 dataset_dir 加载 GNNBenchmarkDataset 数据集。split 参数指定了当前加载的是训练集、验证集还是测试集。
    # [GNNBenchmarkDataset(...) for split in ['train', 'val', 'test']]：生成一个包含训练集、验证集和测试集的列表。
    # join_dataset_splits(...)：将三个数据集（训练、验证、测试）合并为一个单一的 PyG 数据集对象。
    dataset = join_dataset_splits(
        [GNNBenchmarkDataset(root=dataset_dir, name=name, split=split)
         for split in ['train', 'val', 'test']]
    )

    # 使用 pre_transform_in_memory 函数对合并后的数据集应用预变换。
    # pre_transform_in_memory(dataset, T.Compose(tf_list))：对数据集应用变换，变换列表由 tf_list 组成，其中包括 concat_x_and_pos 和 typecast_x。
    # T.Compose(tf_list)：将 tf_list 中的变换函数组合成一个序列，并逐一应用到数据集中的每个样本。
    pre_transform_in_memory(dataset, T.Compose(tf_list))

    return dataset


def preformat_MalNetTiny(dataset_dir, feature_set):
    """Load and preformat Tiny version (5k graphs) of MalNet

    Args:
        dataset_dir: path where to store the cached dataset
        feature_set: select what node features to precompute as MalNet
            originally doesn't have any node nor edge features

    Returns:
        PyG dataset object
    """
    if feature_set in ['none', 'Constant']:
        tf = T.Constant()
    elif feature_set == 'OneHotDegree':
        tf = T.OneHotDegree()
    elif feature_set == 'LocalDegreeProfile':
        tf = T.LocalDegreeProfile()
    else:
        raise ValueError(f"Unexpected transform function: {feature_set}")

    dataset = MalNetTiny(dataset_dir)
    dataset.name = 'MalNetTiny'
    logging.info(f'Computing "{feature_set}" node features for MalNetTiny.')
    pre_transform_in_memory(dataset, tf)

    split_dict = dataset.get_idx_split()
    dataset.split_idxs = [split_dict['train'],
                          split_dict['valid'],
                          split_dict['test']]

    return dataset


def preformat_OGB_Graph(dataset_dir, name):
    """Load and preformat OGB Graph Property Prediction datasets.

    Args:
        dataset_dir: path where to store the cached dataset
        name: name of the specific OGB Graph dataset

    Returns:
        PyG dataset object
    """
    dataset = PygGraphPropPredDataset(name=name, root=dataset_dir)
    s_dict = dataset.get_idx_split()
    dataset.split_idxs = [s_dict[s] for s in ['train', 'valid', 'test']]

    if name == 'ogbg-ppa':
        # ogbg-ppa doesn't have any node features, therefore add zeros but do
        # so dynamically as a 'transform' and not as a cached 'pre-transform'
        # because the dataset is big (~38.5M nodes), already taking ~31GB space
        def add_zeros(data):
            data.x = torch.zeros(data.num_nodes, dtype=torch.long)
            return data
        dataset.transform = add_zeros
    elif name == 'ogbg-code2':
        from graphgps.loader.ogbg_code2_utils import idx2vocab, \
            get_vocab_mapping, augment_edge, encode_y_to_arr
        num_vocab = 5000  # The number of vocabulary used for sequence prediction
        max_seq_len = 5  # The maximum sequence length to predict

        seq_len_list = np.array([len(seq) for seq in dataset.data.y])
        logging.info(f"Target sequences less or equal to {max_seq_len} is "
            f"{np.sum(seq_len_list <= max_seq_len) / len(seq_len_list)}")

        # Building vocabulary for sequence prediction. Only use training data.
        vocab2idx, idx2vocab_local = get_vocab_mapping(
            [dataset.data.y[i] for i in s_dict['train']], num_vocab)
        logging.info(f"Final size of vocabulary is {len(vocab2idx)}")
        idx2vocab.extend(idx2vocab_local)  # Set to global variable to later access in CustomLogger

        # Set the transform function:
        # augment_edge: add next-token edge as well as inverse edges. add edge attributes.
        # encode_y_to_arr: add y_arr to PyG data object, indicating the array repres
        dataset.transform = T.Compose(
            [augment_edge,
             lambda data: encode_y_to_arr(data, vocab2idx, max_seq_len)])

        # Subset graphs to a maximum size (number of nodes) limit.
        pre_transform_in_memory(dataset, partial(clip_graphs_to_size,
                                                 size_limit=1000))

    return dataset


def preformat_OGB_PCQM4Mv2(dataset_dir, name):
    """Load and preformat PCQM4Mv2 from OGB LSC.

    OGB-LSC provides 4 data index splits:
    2 with labeled molecules: 'train', 'valid' meant for training and dev
    2 unlabeled: 'test-dev', 'test-challenge' for the LSC challenge submission

    We will take random 150k from 'train' and make it a validation set and
    use the original 'valid' as our testing set.

    Note: PygPCQM4Mv2Dataset requires rdkit

    Args:
        dataset_dir: path where to store the cached dataset
        name: select 'subset' or 'full' version of the training set

    Returns:
        PyG dataset object
    """
    try:
        # Load locally to avoid RDKit dependency until necessary.
        from ogb.lsc import PygPCQM4Mv2Dataset
    except Exception as e:
        logging.error('ERROR: Failed to import PygPCQM4Mv2Dataset, '
                      'make sure RDKit is installed.')
        raise e


    dataset = PygPCQM4Mv2Dataset(root=dataset_dir)
    split_idx = dataset.get_idx_split()

    rng = default_rng(seed=42)
    train_idx = rng.permutation(split_idx['train'].numpy())
    train_idx = torch.from_numpy(train_idx)

    # Leave out 150k graphs for a new validation set.
    valid_idx, train_idx = train_idx[:150000], train_idx[150000:]
    if name == 'full':
        split_idxs = [train_idx,  # Subset of original 'train'.
                      valid_idx,  # Subset of original 'train' as validation set.
                      split_idx['valid']  # The original 'valid' as testing set.
                      ]
    elif name == 'subset':
        # Further subset the training set for faster debugging.
        subset_ratio = 0.1
        subtrain_idx = train_idx[:int(subset_ratio * len(train_idx))]
        subvalid_idx = valid_idx[:50000]
        subtest_idx = split_idx['valid']  # The original 'valid' as testing set.
        dataset = dataset[torch.cat([subtrain_idx, subvalid_idx, subtest_idx])]
        n1, n2, n3 = len(subtrain_idx), len(subvalid_idx), len(subtest_idx)
        split_idxs = [list(range(n1)),
                      list(range(n1, n1 + n2)),
                      list(range(n1 + n2, n1 + n2 + n3))]
    else:
        raise ValueError(f'Unexpected OGB PCQM4Mv2 subset choice: {name}')
    dataset.split_idxs = split_idxs
    return dataset


def preformat_PCQM4Mv2Contact(dataset_dir, name):
    """Load PCQM4Mv2-derived molecular contact link prediction dataset.

    Note: This dataset requires RDKit dependency!

    Args:
       dataset_dir: path where to store the cached dataset
       name: the type of dataset split: 'shuffle', 'num-atoms'

    Returns:
       PyG dataset object
    """
    try:
        # Load locally to avoid RDKit dependency until necessary
        from graphgps.loader.dataset.pcqm4mv2_contact import \
            PygPCQM4Mv2ContactDataset, \
            structured_neg_sampling_transform
    except Exception as e:
        logging.error('ERROR: Failed to import PygPCQM4Mv2ContactDataset, '
                      'make sure RDKit is installed.')
        raise e

    split_name = name.split('-', 1)[1]
    dataset = PygPCQM4Mv2ContactDataset(dataset_dir, subset='530k')
    # Inductive graph-level split (there is no train/test edge split).
    s_dict = dataset.get_idx_split(split_name)
    dataset.split_idxs = [s_dict[s] for s in ['train', 'val', 'test']]
    if cfg.dataset.resample_negative:
        dataset.transform = structured_neg_sampling_transform
    return dataset


def preformat_Peptides(dataset_dir, name):
    """Load Peptides dataset, functional or structural.

    Note: This dataset requires RDKit dependency!

    Args:
        dataset_dir: path where to store the cached dataset
        name: the type of dataset split:
            - 'peptides-functional' (10-task classification)
            - 'peptides-structural' (11-task regression)

    Returns:
        PyG dataset object
    """
    try:
        # Load locally to avoid RDKit dependency until necessary.
        from graphgps.loader.dataset.peptides_functional import \
            PeptidesFunctionalDataset
        from graphgps.loader.dataset.peptides_structural import \
            PeptidesStructuralDataset
    except Exception as e:
        logging.error('ERROR: Failed to import Peptides dataset class, '
                      'make sure RDKit is installed.')
        raise e

    dataset_type = name.split('-', 1)[1]
    if dataset_type == 'functional':
        dataset = PeptidesFunctionalDataset(dataset_dir)
    elif dataset_type == 'structural':
        dataset = PeptidesStructuralDataset(dataset_dir)
    s_dict = dataset.get_idx_split()
    dataset.split_idxs = [s_dict[s] for s in ['train', 'val', 'test']]
    return dataset


def preformat_TUDataset(dataset_dir, name):
    """Load and preformat datasets from PyG's TUDataset.

    Args:
        dataset_dir: path where to store the cached dataset
        name: name of the specific dataset in the TUDataset class

    Returns:
        PyG dataset object
    """
    if name in ['DD', 'NCI1', 'ENZYMES', 'PROTEINS']:
        func = None
    elif name.startswith('IMDB-') or name == "COLLAB":
        func = T.Constant()
    else:
        ValueError(f"Loading dataset '{name}' from TUDataset is not supported.")
    dataset = TUDataset(dataset_dir, name, pre_transform=func)
    return dataset

def preformat_ogbn(dataset_dir, name):
  if name == 'ogbn-arxiv' or name == 'ogbn-proteins':
    dataset = PygNodePropPredDataset(name=name)
    if name == 'ogbn-arxiv':
      pre_transform_in_memory(dataset, partial(add_reverse_edges))
      if cfg.prep.add_self_loops:
        pre_transform_in_memory(dataset, partial(add_self_loops))
    if name == 'ogbn-proteins':
      pre_transform_in_memory(dataset, partial(move_node_feat_to_x))
      pre_transform_in_memory(dataset, partial(typecast_x, type_str='float'))
    split_dict = dataset.get_idx_split()
    split_dict['val'] = split_dict.pop('valid')
    dataset.split_idx = split_dict
    return dataset


     #  We do not need to store  these separately.
     # storing separatelymight simplify the duplicated logger code in main.py
     # s_dict = dataset.get_idx_split()
     # dataset.split_idxs = [s_dict[s] for s in ['train', 'valid', 'test']]
     # convert the adjacency list to an edge_index list.
     # data = dataset[0]
     # coo = data.adj_t.coo()
     # data is only a deep copy.  Need to write to the dataset object itself.
     # dataset[0].edge_index = torch.stack(coo[:2])
     # del dataset[0]['adj_t'] # remove the adjacency list after the edge_index is created.

     # return dataset
  else:
     ValueError(f"Unknown ogbn dataset '{name}'.")

def preformat_ZINC(dataset_dir, name):
    """Load and preformat ZINC datasets.

    Args:
        dataset_dir: path where to store the cached dataset
        name: select 'subset' or 'full' version of ZINC

    Returns:
        PyG dataset object
    """
    if name not in ['subset', 'full']:
        raise ValueError(f"Unexpected subset choice for ZINC dataset: {name}")
    dataset = join_dataset_splits(
        [ZINC(root=dataset_dir, subset=(name == 'subset'), split=split)
         for split in ['train', 'val', 'test']]
    )
    return dataset


def preformat_AQSOL(dataset_dir):
    """Load and preformat AQSOL datasets.

    Args:
        dataset_dir: path where to store the cached dataset

    Returns:
        PyG dataset object
    """
    dataset = join_dataset_splits(
        [AQSOL(root=dataset_dir, split=split)
         for split in ['train', 'val', 'test']]
    )
    return dataset


def preformat_VOCSuperpixels(dataset_dir, name, slic_compactness):
    """Load and preformat VOCSuperpixels dataset.

    Args:
        dataset_dir: path where to store the cached dataset
    Returns:
        PyG dataset object
    """
    dataset = join_dataset_splits(
        [VOCSuperpixels(root=dataset_dir, name=name,
                        slic_compactness=slic_compactness,
                        split=split)
         for split in ['train', 'val', 'test']]
    )
    return dataset


def preformat_COCOSuperpixels(dataset_dir, name, slic_compactness):
    """Load and preformat COCOSuperpixels dataset.

    Args:
        dataset_dir: path where to store the cached dataset
    Returns:
        PyG dataset object
    """
    dataset = join_dataset_splits(
        [COCOSuperpixels(root=dataset_dir, name=name,
                         slic_compactness=slic_compactness,
                         split=split)
         for split in ['train', 'val', 'test']]
    )
    return dataset


def join_dataset_splits(datasets):
    """Join train, val, test datasets into one dataset object.

    Args:
        datasets: list of 3 PyG datasets to merge

    Returns:
        joint dataset with `split_idxs` property storing the split indices
    """
    # 确保传入的 datasets 列表长度为 3，即需要有且仅有训练集、验证集、测试集三个数据集。如果不是三个数据集，就抛出断言错误，提示 "Expecting train, val, test datasets"。
    assert len(datasets) == 3, "Expecting train, val, test datasets"

    # 获取每个数据集的长度（即包含的样本数量），分别存储在 n1、n2 和 n3 中。其中：
    # n1: 训练集的样本数量
    # n2: 验证集的样本数量
    # n3: 测试集的样本数量
    n1, n2, n3 = len(datasets[0]), len(datasets[1]), len(datasets[2])
    # 使用列表推导式从每个数据集中获取数据样本，通过 datasets[0].get(i) 提取每个样本对象。
    # 通过拼接训练集、验证集、测试集的数据，将它们合并成一个完整的 data_list。
    # 这里 get(i) 是 PyG 数据集对象的方法，用于提取第 i 个样本。
    data_list = [datasets[0].get(i) for i in range(n1)] + \
                [datasets[1].get(i) for i in range(n2)] + \
                [datasets[2].get(i) for i in range(n3)]

    # 清空第一个数据集（即训练集）的 _indices 属性。_indices 是 PyG 中用于存储采样的索引属性，在这里合并后不再需要。
    datasets[0]._indices = None
    # 将新合并的数据列表 data_list 赋值给 datasets[0]._data_list。这一步相当于用合并后的数据替换掉训练集的原始数据列表。
    datasets[0]._data_list = data_list
    # 使用 PyG 的 collate 函数将数据列表 data_list 转换成一个 PyG 数据对象，并根据数据中的批次信息生成 slices（切片）。
    # data 是包含所有样本数据的 PyG 数据对象，slices 用于指示如何分割批次。
    # collate 函数用于将多个图数据拼接在一起，以便处理成一个大图或一个批次的图数据。
    datasets[0].data, datasets[0].slices = datasets[0].collate(data_list)
    # 根据每个数据集的长度创建分割索引 split_idxs，用于指示训练、验证和测试集在合并数据集中的索引范围。具体如下：
    # list(range(n1)): 生成训练集的索引范围 [0, 1, 2, ..., n1-1]。
    # list(range(n1, n1 + n2)): 生成验证集的索引范围 [n1, n1+1, ..., n1+n2-1]。
    # list(range(n1 + n2, n1 + n2 + n3)): 生成测试集的索引范围 [n1+n2, ..., n1+n2+n3-1]。
    split_idxs = [list(range(n1)),
                  list(range(n1, n1 + n2)),
                  list(range(n1 + n2, n1 + n2 + n3))]
    datasets[0].split_idxs = split_idxs

    return datasets[0]